#ifndef AVL_TREE_H
#define AVL_TREE_H

#include <iostream>

template <typename T, typename Compare = std::less<T>>
class AVLTree
{
private:
    struct Node
    {
        Node(T val = T())
            : val_(val), left_(nullptr), right_(nullptr), height_(1)
        {
        }

        T val_;
        Node *left_;
        Node *right_;
        int height_;
    };

    Node *root_;
    Compare comp_;

private:
    // 获取当前节点的高度
    int height(Node *node) const
    {
        if (!node)
            return 0;
        return node->height_;
    }

    // 右旋转
    Node *rightRotate(Node *node) const
    {
        Node *child = node->left_;
        node->left_ = child->right_;
        child->right_ = node;

        node->height_ = std::max(height(node->left_), height(node->right_)) + 1;
        child->height_ = std::max(height(child->left_), height(child->right_)) + 1;

        return child;
    }

    // 左旋转
    Node *leftRotate(Node *node) const
    {
        Node *child = node->right_;
        node->right_ = child->left_;
        child->left_ = node;

        node->height_ = std::max(height(node->left_), height(node->right_)) + 1;
        child->height_ = std::max(height(child->left_), height(child->right_)) + 1;

        return child;
    }

    // 右平衡
    Node *rightBalance(Node *node) const
    {
        Node *child = node->right_;
        node->right_ = rightRotate(child);
        return leftRotate(node);
    }

    // 左平衡
    Node *leftBalance(Node *node) const
    {
        Node *child = node->left_;
        node->left_ = leftRotate(child);
        return rightRotate(node);
    }

    Node *insert(Node *node, const T &val)
    {
        if (!node)
        {
            return new Node(val);
        }
        if (!comp_(node->val_, val))
        {
            node->left_ = insert(node->left_, val);
            // 在递归回溯的过程中，如果当前节点失去平衡
            if (height(node->left_) - height(node->right_) > 1)
            {
                // 如果孩子节点的左子树高于右子树,就进行左旋，否则进行左平衡
                if (height(node->left_->left_) >= height(node->left_->right_))
                {
                    node = rightRotate(node);
                }
                else
                {
                    node = leftBalance(node);
                }
            }
        }
        else if (comp_(node->val_, val))
        {
            node->right_ = insert(node->right_, val);
            // 如果孩子节点的右子树高于左子树,就进行左旋，否则进行平衡
            if (height(node->right_) - height(node->left_) > 1)
            {
                // 如果当前
                if (height(node->right_->left_) >= height(node->right_->right_))
                {
                    node = rightBalance(node);
                }
                else
                {
                    node = leftRotate(node);
                }
            }
        }
        node->height_ = std::max(height(node->left_), height(node->right_)) + 1;
        return node;
    }

    Node *erase(Node *node, const T &val)
    {
        // 如果当前节点不存在，直接返回
        if (!node)
            return nullptr;
        if (node->val_ == val)
        {
            // 如果当前节点的左右孩子都不为nullptr
            if (node->left_ && node->right_)
            {
                // 删除当前左右子树中高度更高的节点
                if(height(node->left_) >= height(node->right_)){
                    Node *pre = node->left_;
                    while(pre->right_){
                        pre = pre->right_;
                    }
                    node->val_ = pre->val_;
                    node->left_ = erase(node->left_, pre->val_);
                }else{
                    Node *suf = node->right_;
                    while(suf->left_){
                        suf = suf->left_;
                    }
                    node->val_ = suf->val_;
                    node->right_ = erase(node->right_, suf->val_);
                }
            }
            else
            {
                Node *child = node->left_;
                if (!child)
                    child = node->right_;
                delete node;
                return child;
            }
        }
        else if (!comp_(node->val_, val))
        {
            node->left_ = erase(node->left_, val);
            // 在递归回溯的过程中，如果当前节点失去平衡
            if (height(node->right_) - height(node->left_) > 1)
            {
                // 如果孩子节点的左子树高于右子树,就进行左旋，否则进行左平衡
                if (height(node->right_->left_) >= height(node->right_->right_))
                {
                    node = rightBalance(node);
                }
                else
                {
                    node = leftRotate(node);
                }
            }
        }
        else
        {
            node->right_ = erase(node->right_, val);
            // 如果孩子节点的右子树高于左子树,就进行左旋，否则进行平衡
            if (height(node->left_) - height(node->right_) > 1)
            {
                // 如果当前
                if (height(node->left_->left_) >= height(node->left_->right_))
                {
                    node = rightRotate(node);
                }
                else
                {
                    node = leftBalance(node);
                }
            }
        }
        node->height_ = std::max(height(node->left_), height(node->right_)) + 1;
        return node;
    }

public:
    AVLTree(Compare comp = Compare()) : root_(nullptr), comp_(comp) {}
    ~AVLTree() {}

    void erase(const T &val)
    {
        if (!root_)
            return;
        root_ = erase(root_, val);
    }

    void insert(const T &val)
    {
        root_ = insert(root_, val);
    }
};

#endif